/**
 * 
 */
package edu.umd.clip.lm.programs;

import java.io.IOException;
import java.lang.management.*;
import java.net.*;
import java.util.*;

import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;

import edu.berkeley.nlp.util.*;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.InterpolatingProbTreeStorage;
import edu.umd.clip.lm.storage.*;
import edu.umd.clip.lm.storage.remote.*;
import edu.umd.clip.lm.storage.remote.ServerStorageConnection.Task;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LMStorageServer {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-forest", required = true, usage = "the decision tree forest")
		public String forest;        
        @Option(name = "-port", required = false, usage = "TCP port to listen at (default: 2332)")
        public int port = 2332;
        @Option(name = "-meminfo", usage = "track memory usage (default: false)")
        public boolean meminfo = false;
	}

	/**
	 * @param args
	 * @throws DatabaseException 
	 * @throws IOException 
	 */
	public static void main(String[] args) throws DatabaseException, IOException {
		final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
		
		MemoryUsage memuse;
		
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

        SocketAddress address = new InetSocketAddress(opts.port);
		final ServerStorageConnection connection = new ServerStorageConnection(address);

		Experiment.initialize(opts.config);
		final Experiment experiment = Experiment.getInstance();

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after Experiment.initialize(): %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}
		
		experiment.buildPrefixes();
		experiment.buildWordPrefixes();

		DecoderCompactProbTree emptyTree = DecoderCompactProbTree.constructEmptyTree(Experiment.getInstance().getHFT().getTree());
		DecoderCompactProbTree.setEmptyTree(emptyTree);
		
		Experiment.Files files = experiment.getFiles();
		Environment env = BDBProbTreeStorage.createEnvironment(files.getDb(), false);
		env.getConfig().setCachePercent(50);
		
		final ForestModel forest = experiment.getForest(opts.forest);
		final ServerStorageResponder responder;
		{
			BDBProbTreeStorage dbStorage = new BDBProbTreeStorage(env);
			InterpolatingProbTreeStorage interpolatingStorage = new InterpolatingProbTreeStorage(dbStorage);
			interpolatingStorage.setCacheSize(300000);
			responder = new ServerStorageResponder(interpolatingStorage, connection);
			// load the chain of LMs
			for(LanguageModel lm : forest.getModels()) {
				System.err.printf("loading %s history tree\n", lm.getId());
				lm.getHistoryTree();
				System.err.printf("loading %s db\n", lm.getId());
				dbStorage.open(lm.getId(), lm.getIdNum(), false);
				interpolatingStorage.addLm(lm);
			}
			forest.getDecoder().setStorage(interpolatingStorage);
		}

		for(int i=0; i<opts.jobs; ++i) {
			new Worker(connection, responder).start();
		}
		
		Timer timer = new Timer(true);
		timer.scheduleAtFixedRate(new ServerStorageConnection.ServerMonitor(connection), 0, ServerStorageConnection.ServerMonitor.INTERVAL);
		System.err.println("Ready to serve...");
		//connection.serverLoop();
		try {
			while(true) {
				Thread.sleep(1000000);
			}
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
	}

	private static class Worker extends Thread {
		ServerStorageConnection connection;
		ServerStorageResponder responder;
		
		public Worker(ServerStorageConnection connection, ServerStorageResponder responder) {
			this.connection = connection;
			this.responder = responder;
		}

		public void run() {
			while(!isInterrupted()) {
				Task task = connection.getNextTask();
				if (task != null) {
					responder.process(task);
				}
			}
		}
	}
}
